import json
import uuid

from conf.neo4j_cnx import NeoCnxPool
from medical_kg.med_kg import MedKg
from medical_kg.prop_set_up import prop_model

neo_cnx = NeoCnxPool()


class KgAlter:

    def query_rel_type_list(self, start_type_name):
        with neo_cnx.get_session() as session:
            rel_type_list = []

            query = (
                'MATCH (n1:EntityType) '
                'MATCH (n2:EntityType) '
                'MATCH (n1)-[rel]->(n2) '
                'WHERE n1.name = $start_type_name AND rel.is_abled = True '
                'RETURN n1, n2, rel '
                'ORDER BY rel.update_time DESC '
            )
            results = session.run(query, start_type_name=start_type_name)

            for record in results:
                for i in ['rel']:
                    rel_type = record[i].type
                    end_name = record[i].nodes[1].get('name')

                    rel_type_list.append({
                        'rel_type': rel_type,
                        'end_name': end_name
                    })

            return rel_type_list

    def add_node_proerties(self, uid, properties: dict):
        with neo_cnx.get_session() as session:
            query = (
                'MATCH (n:ICD11 {uid:$uid}) '
                'RETURN n.properties AS properties'
            )
            results = session.run(query, uid=uid)
            for record in results:
                old = json.loads(record['properties'])

                new = properties
                if 'name' in new: new.pop('name')
                if 'code' in new: new.pop('code')

                diff = {k: old[k] for k in set(old) - set(new)}
                z = {**diff, **new}

                properties = json.dumps(z, ensure_ascii=False)
                query = (
                    'MATCH (n:ICD11 {uid:$uid}) '
                    'SET n.properties=$properties'
                )
                session.run(query, uid=uid, properties=properties)

    def create_rel(self, uid1, uid2, rel_type, rel_properties: dict):
        with neo_cnx.get_session() as session:
            rel_properties = json.dumps(prop_model(rel_properties, type='REL'), ensure_ascii=False)
            query = (
                        'MATCH (n1:ICD11 {uid:$uid1}) '
                        'MATCH (n2:ICD11 {uid:$uid2}) '
                        'MERGE (n1)-[r:%s]->(n2) '
                        'ON CREATE SET r.properties=$rel_properties, r.uid=$rel_uid'
                    ) % rel_type
            session.run(query, uid1=uid1, uid2=uid2, rel_properties=rel_properties, rel_uid=str(uuid.uuid4()))

    def del_rel(self, uid1, uid2):
        with neo_cnx.get_session() as session:
            query = (
                'MATCH (n1:ICD11 {uid:$uid1}) '
                'MATCH (n2:ICD11 {uid:$uid2}) '
                'MATCH (n1)-[r]-(n2) '
                'RETURN r'
            )
            results = session.run(query, uid1=uid1, uid2=uid2)
            for record in results:
                rel_type = record['r'].type
                if rel_type != 'Parent' and rel_type != 'Child':
                    glabal_properties = dict(record['r'].items())
                    rel_uid = glabal_properties['uid']
                    query = (
                        'MATCH ()-[r]-() WHERE r.uid=%rel_uid DELETE r'
                    )
                    session.run(query, rel_uid=rel_uid)

    def alter_gragh(self, uid, gragh: dict):
        node_list = gragh['nodes']
        for node in node_list:
            uid = node['id']
            properties = node['properties']
            self.add_node_proerties(uid, prop_model(properties))

        new_rel_list = gragh['relationships']

        rel_type_list = []
        for rel in new_rel_list:
            rel_type_list.append(rel['type'])

        if 'Parent' in rel_type_list: rel_type_list.remove('Parent')
        if 'Child' in rel_type_list: rel_type_list.remove('Child')

        # delete the current relationships
        cur_rel_list = MedKg().query_disease_kg(uid, rel_type_list)['graph']['relationships']
        for rel in cur_rel_list:
            self.del_rel(rel['startNode'], rel['endNode'])

        # create the new relathionships
        for rel in new_rel_list:
            self.create_rel(rel['startNode'], rel['endNode'], rel['type'], rel['properties'])


if __name__ == '__main__':
    KgAlter().add_node_proerties(uid='1435254666', properties=prop_model({'名称': '某些感染性疾病或寄生虫病',
                                                                          '编码': '01'}))
